"""
文件IO操作工具模块

提供文件读写、路径管理、数据序列化等功能
"""

import os
import json
import pickle
import yaml
from pathlib import Path
from typing import Any, Dict, List, Optional, Union, Tuple
import numpy as np
import torch
from PIL import Image
import cv2
import pandas as pd


def ensure_dir(path: Union[str, Path]) -> Path:
    """确保目录存在，如果不存在则创建"""
    path = Path(path)
    path.mkdir(parents=True, exist_ok=True)
    return path


def load_json(file_path: Union[str, Path]) -> Dict[str, Any]:
    """加载JSON文件"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return json.load(f)


def save_json(data: Dict[str, Any], file_path: Union[str, Path], indent: int = 2) -> None:
    """保存数据为JSON文件"""
    ensure_dir(Path(file_path).parent)
    with open(file_path, 'w', encoding='utf-8') as f:
        json.dump(data, f, indent=indent, ensure_ascii=False)


def load_yaml(file_path: Union[str, Path]) -> Dict[str, Any]:
    """加载YAML配置文件"""
    with open(file_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)


def save_yaml(data: Dict[str, Any], file_path: Union[str, Path]) -> None:
    """保存数据为YAML文件"""
    ensure_dir(Path(file_path).parent)
    with open(file_path, 'w', encoding='utf-8') as f:
        yaml.dump(data, f, default_flow_style=False, allow_unicode=True)


def load_pickle(file_path: Union[str, Path]) -> Any:
    """加载pickle文件"""
    with open(file_path, 'rb') as f:
        return pickle.load(f)


def save_pickle(data: Any, file_path: Union[str, Path]) -> None:
    """保存数据为pickle文件"""
    ensure_dir(Path(file_path).parent)
    with open(file_path, 'wb') as f:
        pickle.dump(data, f)


def load_image(file_path: Union[str, Path], 
               mode: str = 'RGB',
               backend: str = 'auto',
               target_size: Optional[Tuple[int, int]] = None) -> np.ndarray:
    """
    加载图像文件
    
    Args:
        file_path: 图像文件路径
        mode: 图像模式 ('RGB', 'GRAY', 'RGBA', 'BGR')
        backend: 后端库 ('auto', 'cv2', 'PIL')
        target_size: 目标尺寸 (width, height)，如果指定则自动调整大小
        
    Returns:
        numpy数组格式的图像数据
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"图像文件不存在: {file_path}")
    
    # 自动选择后端
    if backend == 'auto':
        ext = file_path.suffix.lower()
        # 对于某些格式，PIL支持更好
        if ext in ['.tiff', '.tif', '.webp', '.ico']:
            backend = 'PIL'
        else:
            backend = 'cv2'
    
    if backend == 'PIL':
        # 使用PIL加载
        try:
            pil_image = Image.open(file_path)
            
            # 转换为RGB模式（如果需要）
            if mode == 'RGB' and pil_image.mode != 'RGB':
                if pil_image.mode == 'RGBA':
                    # 创建白色背景
                    background = Image.new('RGB', pil_image.size, (255, 255, 255))
                    background.paste(pil_image, mask=pil_image.split()[3] if pil_image.mode == 'RGBA' else None)
                    pil_image = background
                else:
                    pil_image = pil_image.convert('RGB')
            elif mode == 'GRAY' and pil_image.mode != 'L':
                pil_image = pil_image.convert('L')
            elif mode == 'RGBA' and pil_image.mode != 'RGBA':
                pil_image = pil_image.convert('RGBA')
            
            # 调整大小
            if target_size is not None:
                pil_image = pil_image.resize(target_size, Image.Resampling.LANCZOS)
            
            # 转换为numpy数组
            image = np.array(pil_image)
            
            # 如果是灰度图，确保是2D数组
            if mode == 'GRAY' and len(image.shape) == 3:
                image = image[:, :, 0]
                
        except Exception as e:
            raise ValueError(f"使用PIL加载图像失败 {file_path}: {e}")
    else:
        # 使用OpenCV加载
        if mode == 'GRAY':
            image = cv2.imread(str(file_path), cv2.IMREAD_GRAYSCALE)
        elif mode == 'RGBA':
            image = cv2.imread(str(file_path), cv2.IMREAD_UNCHANGED)
            if image is not None and len(image.shape) == 3 and image.shape[2] == 3:
                # 添加alpha通道
                image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)
        else:
            image = cv2.imread(str(file_path))
        
        if image is None:
            raise ValueError(f"无法加载图像: {file_path}")
        
        # 颜色空间转换
        if mode == 'RGB' and len(image.shape) == 3:
            if image.shape[2] == 4:  # BGRA
                image = cv2.cvtColor(image, cv2.COLOR_BGRA2RGB)
            else:  # BGR
                image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
        elif mode == 'RGBA' and len(image.shape) == 3 and image.shape[2] == 3:
            image = cv2.cvtColor(image, cv2.COLOR_BGR2BGRA)
        
        # 调整大小
        if target_size is not None:
            if mode == 'GRAY':
                image = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
            else:
                image = cv2.resize(image, target_size, interpolation=cv2.INTER_LINEAR)
    
    return image


def save_image(image: np.ndarray, file_path: Union[str, Path]) -> None:
    """保存图像文件"""
    ensure_dir(Path(file_path).parent)
    if len(image.shape) == 3 and image.shape[2] == 3:
        image = cv2.cvtColor(image, cv2.COLOR_RGB2BGR)
    cv2.imwrite(str(file_path), image)


def load_torch_model(file_path: Union[str, Path], device: str = 'cpu') -> torch.nn.Module:
    """加载PyTorch模型"""
    return torch.load(file_path, map_location=device)


def save_torch_model(model: torch.nn.Module, file_path: Union[str, Path]) -> None:
    """保存PyTorch模型"""
    ensure_dir(Path(file_path).parent)
    torch.save(model.state_dict(), file_path)


def get_file_list(directory: Union[str, Path], 
                  extensions: Optional[List[str]] = None,
                  recursive: bool = True) -> List[Path]:
    """获取目录下的文件列表"""
    directory = Path(directory)
    if not directory.exists():
        return []
    
    if extensions is None:
        extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.npy', '.npz']
    
    pattern = "**/*" if recursive else "*"
    files = []
    for ext in extensions:
        files.extend(directory.glob(f"{pattern}{ext}"))
        files.extend(directory.glob(f"{pattern}{ext.upper()}"))
    
    return sorted(files)


def get_subdirs(directory: Union[str, Path]) -> List[Path]:
    """获取目录下的所有子目录"""
    directory = Path(directory)
    if not directory.exists():
        return []
    
    return [d for d in directory.iterdir() if d.is_dir()]


def copy_file(src: Union[str, Path], dst: Union[str, Path]) -> None:
    """复制文件"""
    import shutil
    ensure_dir(Path(dst).parent)
    shutil.copy2(src, dst)


def move_file(src: Union[str, Path], dst: Union[str, Path]) -> None:
    """移动文件"""
    import shutil
    ensure_dir(Path(dst).parent)
    shutil.move(src, dst)


def get_file_size(file_path: Union[str, Path]) -> int:
    """获取文件大小（字节）"""
    return Path(file_path).stat().st_size


def get_file_extension(file_path: Union[str, Path]) -> str:
    """获取文件扩展名"""
    return Path(file_path).suffix.lower()


def is_image_file(file_path: Union[str, Path]) -> bool:
    """判断是否为图像文件"""
    image_extensions = {'.jpg', '.jpeg', '.png', '.bmp', '.tiff', '.tif'}
    return get_file_extension(file_path) in image_extensions


def is_video_file(file_path: Union[str, Path]) -> bool:
    """判断是否为视频文件"""
    video_extensions = {'.mp4', '.avi', '.mov', '.mkv', '.wmv', '.flv'}
    return get_file_extension(file_path) in video_extensions


# ==================== Excel读取工具 ====================

def load_excel(file_path: Union[str, Path],
               sheet_name: Optional[Union[str, int, List[Union[str, int]]]] = None,
               header: Optional[int] = 0,
               skiprows: Optional[int] = None,
               nrows: Optional[int] = None,
               usecols: Optional[Union[str, List[int], List[str]]] = None,
               engine: Optional[str] = None) -> Union[pd.DataFrame, Dict[str, pd.DataFrame]]:
    """
    加载Excel文件
    
    Args:
        file_path: Excel文件路径 (.xlsx, .xls)
        sheet_name: 工作表名称或索引，可以是字符串、整数或列表。None表示加载所有工作表
        header: 表头行号，默认为0（第一行）
        skiprows: 跳过的行数
        nrows: 读取的行数限制
        usecols: 要读取的列，可以是列名列表、列索引列表或字符串（如'A:C'）
        engine: 引擎 ('openpyxl' for .xlsx, 'xlrd' for .xls, None自动选择)
        
    Returns:
        如果sheet_name是单个值，返回DataFrame；如果是列表或None，返回字典 {sheet_name: DataFrame}
        
    Examples:
        # 读取单个工作表
        df = load_excel('data.xlsx', sheet_name='Sheet1')
        
        # 读取所有工作表
        sheets = load_excel('data.xlsx')  # 返回字典
        
        # 读取指定列
        df = load_excel('data.xlsx', usecols=['A', 'B', 'C'])
        
        # 读取前100行
        df = load_excel('data.xlsx', nrows=100)
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"Excel文件不存在: {file_path}")
    
    # 自动选择引擎
    if engine is None:
        ext = file_path.suffix.lower()
        if ext == '.xlsx':
            try:
                import openpyxl
                engine = 'openpyxl'
            except ImportError:
                raise ImportError("需要安装openpyxl来读取.xlsx文件: pip install openpyxl")
        elif ext == '.xls':
            try:
                import xlrd
                engine = 'xlrd'
            except ImportError:
                raise ImportError("需要安装xlrd来读取.xls文件: pip install xlrd")
        else:
            raise ValueError(f"不支持的Excel文件格式: {ext}")
    
    # 读取参数
    read_kwargs = {
        'header': header,
        'engine': engine
    }
    if skiprows is not None:
        read_kwargs['skiprows'] = skiprows
    if nrows is not None:
        read_kwargs['nrows'] = nrows
    if usecols is not None:
        read_kwargs['usecols'] = usecols
    
    # 读取数据
    if sheet_name is None:
        # 读取所有工作表
        excel_file = pd.ExcelFile(file_path, engine=engine)
        return {sheet: pd.read_excel(excel_file, sheet_name=sheet, **read_kwargs) 
                for sheet in excel_file.sheet_names}
    elif isinstance(sheet_name, list):
        # 读取多个指定工作表
        excel_file = pd.ExcelFile(file_path, engine=engine)
        return {sheet: pd.read_excel(excel_file, sheet_name=sheet, **read_kwargs) 
                for sheet in sheet_name if sheet in excel_file.sheet_names}
    else:
        # 读取单个工作表
        return pd.read_excel(file_path, sheet_name=sheet_name, **read_kwargs)


def get_excel_sheet_names(file_path: Union[str, Path]) -> List[str]:
    """
    获取Excel文件中的所有工作表名称
    
    Args:
        file_path: Excel文件路径
        
    Returns:
        工作表名称列表
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"Excel文件不存在: {file_path}")
    
    ext = file_path.suffix.lower()
    if ext == '.xlsx':
        try:
            import openpyxl
            engine = 'openpyxl'
        except ImportError:
            raise ImportError("需要安装openpyxl来读取.xlsx文件: pip install openpyxl")
    elif ext == '.xls':
        try:
            import xlrd
            engine = 'xlrd'
        except ImportError:
            raise ImportError("需要安装xlrd来读取.xls文件: pip install xlrd")
    else:
        raise ValueError(f"不支持的Excel文件格式: {ext}")
    
    excel_file = pd.ExcelFile(file_path, engine=engine)
    return excel_file.sheet_names


def save_excel(data: Union[pd.DataFrame, Dict[str, pd.DataFrame]],
               file_path: Union[str, Path],
               sheet_name: Optional[Union[str, List[str]]] = None,
               index: bool = False,
               engine: Optional[str] = None) -> None:
    """
    保存数据到Excel文件
    
    Args:
        data: 要保存的数据，可以是DataFrame或字典 {sheet_name: DataFrame}
        file_path: 保存路径
        sheet_name: 工作表名称（当data是DataFrame时使用）或名称列表
        index: 是否保存索引
        engine: 引擎 ('openpyxl' for .xlsx, None自动选择)
        
    Examples:
        # 保存单个DataFrame
        save_excel(df, 'output.xlsx', sheet_name='Data')
        
        # 保存多个DataFrame
        save_excel({'Sheet1': df1, 'Sheet2': df2}, 'output.xlsx')
    """
    file_path = Path(file_path)
    ensure_dir(file_path.parent)
    
    # 自动选择引擎
    if engine is None:
        ext = file_path.suffix.lower()
        if ext == '.xlsx':
            try:
                import openpyxl
                engine = 'openpyxl'
            except ImportError:
                raise ImportError("需要安装openpyxl来保存.xlsx文件: pip install openpyxl")
        else:
            raise ValueError(f"不支持的Excel文件格式: {ext}")
    
    if isinstance(data, pd.DataFrame):
        # 单个DataFrame
        if sheet_name is None:
            sheet_name = 'Sheet1'
        data.to_excel(file_path, sheet_name=sheet_name, index=index, engine=engine)
    elif isinstance(data, dict):
        # 多个DataFrame
        with pd.ExcelWriter(file_path, engine=engine) as writer:
            for sheet, df in data.items():
                df.to_excel(writer, sheet_name=sheet, index=index)
    else:
        raise TypeError(f"不支持的数据类型: {type(data)}")


# ==================== 图片工具增强 ====================

def load_image_pil(file_path: Union[str, Path],
                   mode: str = 'RGB',
                   target_size: Optional[Tuple[int, int]] = None) -> Image.Image:
    """
    使用PIL加载图像文件（返回PIL Image对象）
    
    Args:
        file_path: 图像文件路径
        mode: 图像模式 ('RGB', 'RGBA', 'L'等)
        target_size: 目标尺寸 (width, height)
        
    Returns:
        PIL Image对象
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"图像文件不存在: {file_path}")
    
    image = Image.open(file_path)
    
    # 转换模式
    if mode and image.mode != mode:
        if mode == 'RGB' and image.mode == 'RGBA':
            # 创建白色背景
            background = Image.new('RGB', image.size, (255, 255, 255))
            background.paste(image, mask=image.split()[3] if image.mode == 'RGBA' else None)
            image = background
        else:
            image = image.convert(mode)
    
    # 调整大小
    if target_size is not None:
        image = image.resize(target_size, Image.Resampling.LANCZOS)
    
    return image


def get_image_info(file_path: Union[str, Path]) -> Dict[str, Any]:
    """
    获取图像文件信息
    
    Args:
        file_path: 图像文件路径
        
    Returns:
        包含图像信息的字典
    """
    file_path = Path(file_path)
    if not file_path.exists():
        raise FileNotFoundError(f"图像文件不存在: {file_path}")
    
    try:
        with Image.open(file_path) as img:
            info = {
                'width': img.width,
                'height': img.height,
                'mode': img.mode,
                'format': img.format,
                'size_bytes': file_path.stat().st_size,
                'has_transparency': img.mode in ('RGBA', 'LA') or 'transparency' in img.info
            }
            return info
    except Exception as e:
        raise ValueError(f"无法读取图像信息 {file_path}: {e}")


def is_excel_file(file_path: Union[str, Path]) -> bool:
    """判断是否为Excel文件"""
    excel_extensions = {'.xlsx', '.xls', '.xlsm'}
    return get_file_extension(file_path) in excel_extensions

